#!/usr/bin/env python3
"""
Compare Fortran GSI fort.207 diagnostics with Julia radiance summary statistics.
"""
import argparse
import csv
import json
import math
from pathlib import Path
from typing import Dict, List


def parse_fort207(path: Path) -> List[Dict[str, float]]:
    rows: List[Dict[str, float]] = []
    with path.open() as f:
        for raw in f:
            line = raw.strip()
            if not line or line.startswith("RADINFO"):
                continue
            if not line[0].isdigit() or "chan=" not in line:
                continue
            tokens = line.split()
            entry: Dict[str, float] = {
                "index": float(tokens[0]),
                "instrument": tokens[1],
            }
            for key_token, value_token in zip(tokens[2::2], tokens[3::2]):
                if not key_token.endswith("="):
                    continue
                key = key_token[:-1]
                if key in {"chan", "use", "icld_det", "icloud", "iaeros"}:
                    entry[key] = int(float(value_token))
                else:
                    entry[key] = float(value_token)
            rows.append(entry)
    return rows


def load_julia_summary(path: Path) -> Dict[int, Dict[str, float]]:
    data = json.loads(path.read_text())
    stats = data.get("channel_statistics", {})
    channel_data: Dict[int, Dict[str, float]] = {}
    for ch_str, values in stats.items():
        ch = int(ch_str)
        channel_data[ch] = {
            "julia_mean_tb": values.get("mean_tb"),
            "julia_std_tb": values.get("std_tb"),
            "julia_min_tb": values.get("min_tb"),
            "julia_max_tb": values.get("max_tb"),
            "julia_error_std": values.get("obs_error_std"),
            "julia_error_var": values.get("obs_error_var"),
            "julia_use_flag": values.get("use_flag"),
            "julia_valid_fraction": values.get("valid_fraction"),
            "julia_valid_count": values.get("valid_count"),
        }
    return channel_data


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("fort207", type=Path, help="Path to Fortran fort.207 file")
    parser.add_argument(
        "summary_json",
        type=Path,
        help="Path to Julia radiance summary JSON (results/radiance_real_summary.json)",
    )
    parser.add_argument(
        "output_csv",
        type=Path,
        help="Output CSV path for merged comparison",
    )
    parser.add_argument(
        "--instrument",
        type=str,
        default="amsua_n15",
        help="Instrument key to compare (default: amsua_n15)",
    )
    args = parser.parse_args()

    fort_rows = [row for row in parse_fort207(args.fort207) if row.get("instrument") == args.instrument]
    if not fort_rows:
        raise SystemExit(f"Instrument {args.instrument} not found in {args.fort207}")

    julia_stats = load_julia_summary(args.summary_json)

    fort_rows.sort(key=lambda r: int(r.get("chan", 0)))

    args.output_csv.parent.mkdir(parents=True, exist_ok=True)
    fieldnames = [
        "chan",
        "var",
        "varch_cld",
        "use",
        "ermax",
        "b_rad",
        "pg_rad",
        "icld_det",
        "icloud",
        "iaeros",
        "julia_mean_tb",
        "julia_std_tb",
        "julia_min_tb",
        "julia_max_tb",
        "julia_error_std",
        "julia_error_var",
        "julia_use_flag",
        "julia_valid_fraction",
        "julia_valid_count",
        "error_std_diff",
        "error_std_rel_pct",
        "error_var_diff",
        "error_var_rel_pct",
    ]

    with args.output_csv.open("w", newline="") as csvfile:
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        writer.writeheader()
        for row in fort_rows:
            chan = int(row.get("chan", 0))
            merged = {key: row.get(key) for key in fieldnames}
            merged.update(julia_stats.get(chan, {}))

            var_value = row.get("var")
            julia_var = merged.get("julia_error_var")
            julia_std = merged.get("julia_error_std")
            if var_value not in (None, "") and julia_std not in (None, ""):
                std_diff = float(julia_std) - float(var_value)
                merged["error_std_diff"] = std_diff
                merged["error_std_rel_pct"] = std_diff / float(var_value) * 100.0 if float(var_value) != 0.0 else None
            else:
                merged["error_std_diff"] = None
                merged["error_std_rel_pct"] = None

            if var_value not in (None, "") and julia_var not in (None, ""):
                diff = float(julia_var) - float(var_value)
                merged["error_var_diff"] = diff
                merged["error_var_rel_pct"] = diff / float(var_value) * 100.0 if float(var_value) != 0.0 else None
            else:
                merged["error_std_diff"] = merged.get("error_std_diff")
                merged["error_std_rel_pct"] = merged.get("error_std_rel_pct")
                merged["error_var_diff"] = None
                merged["error_var_rel_pct"] = None
            writer.writerow(merged)

    print(f"Wrote comparison CSV to {args.output_csv}")

    md_path = args.output_csv.with_suffix(".md")
    with md_path.open("w") as md:
        md.write(f"# Fortran fort.207 vs Julia radiance summary ({args.instrument})\n\n")
        md.write(f"Source Fortran file: `{args.fort207}`\\n\n")
        md.write("| Channel | Fortran σ (K) | Fortran use | Julia σ (K) | Δσ (K) | Δσ (%) | Julia use | Valid % | Julia mean TB (K) | Julia std (K) |\n")
        md.write("|--------:|---------------:|------------:|------------:|---------:|--------:|----------:|--------:|------------------:|--------------:|\n")
        for row in fort_rows:
            chan = int(row.get("chan", 0))
            stats = julia_stats.get(chan, {})
            var = row.get("var", float("nan"))
            julia_std = stats.get("julia_error_std", float("nan")) or float("nan")
            julia_var = stats.get("julia_error_var", float("nan")) or float("nan")
            std_diff = julia_std - var if not (math.isnan(julia_std) or math.isnan(var)) else float("nan")
            rel_std = std_diff / var * 100.0 if not (math.isnan(std_diff) or var == 0) else float("nan")
            use_flag = row.get("use", 0)
            julia_use = stats.get("julia_use_flag", float("nan")) or float("nan")
            valid_pct = stats.get("julia_valid_fraction", float("nan"))
            if valid_pct is not None and not math.isnan(valid_pct):
                valid_pct *= 100.0
            else:
                valid_pct = float("nan")
            md.write(
                "| {chan:2d} | {var:6.3f} | {use:6.0f} | {sigma:10.3f} | {ds:7.3f} | {rel:6.2f} | {juse:8.0f} | {valid:7.2f} | {mean:8.3f} | {std:6.3f} |\n".format(
                    chan=chan,
                    var=var,
                    use=float(use_flag),
                    sigma=julia_std,
                    ds=std_diff,
                    rel=rel_std,
                    juse=julia_use,
                    valid=valid_pct,
                    mean=stats.get("julia_mean_tb", float("nan")) or float("nan"),
                    std=stats.get("julia_std_tb", float("nan")) or float("nan"),
                )
            )
    print(f"Wrote Markdown summary to {md_path}")


if __name__ == "__main__":
    main()
